{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/langroid/langroid/blob/main/examples/langroid_quick_examples.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uIV7QkkrC8O7"
   },
   "source": [
    "\n",
    "\n",
    "<img width=\"700\" src=\"https://raw.githubusercontent.com/langroid/langroid/main/docs/assets/langroid-card-lambda-ossem-rust-1200-630.png\" alt=\"Langroid\">\n",
    "\n",
    "# Overview\n",
    "\n",
    "This notebook provides the runnable code for the six [**Usage Examples**](https://github.com/langroid/langroid#tada-usage-examples) described in [Langroid repo](https://github.com/langroid/langroid).\n",
    "\n",
    "**NOTE:** Notebooks (colab, jupyter, or otherwise) are *not* an ideal way to run interactive chat loops. We are showing these examples here since we recognize that Colab notebooks offer the benefit of having a ready to run environment with minimal setup. But we encourage you to try the python scripts in the [examples folder](https://github.com/langroid/langroid/tree/main/examples) of the repo on the command line for the best experience.\n",
    "\n",
    "In the first two cells we show the steps for setting up the requirements to run the examples including the installation of `Langroid` package and setting up the `OPENAI_API_KEY`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hoTp_cNcriIg"
   },
   "source": [
    "## Install Langroid\n",
    "\n",
    "At the end there may be a message saying \"RESTART RUNTIME\", which can be safely ignored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PYaFworprwEJ"
   },
   "outputs": [],
   "source": [
    "!pip install langroid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "v-BiRXu9JQ5H"
   },
   "source": [
    "## Set up `OPENAI_API_KEY`\n",
    "\n",
    "This code will ask the user to provide the `OPENAI_API_KEY`. Before running this cell, please follow these steps to get the key.\n",
    "Login to your OpenAI account --> go to `View API Keys` from the drop-down list on the top-right corner --> click on the botton **create new secret key** --> a new screen will pop up --> press the botton **create secret key**.\n",
    "\n",
    "Visit [this page](https://help.openai.com/en/articles/4936850-where-do-i-find-my-secret-api-key) for more info about where to find the API Key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GOR8OsfvsN2k"
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "\n",
    "import nest_asyncio\n",
    "\n",
    "logging.getLogger().setLevel(logging.ERROR)\n",
    "nest_asyncio.apply()\n",
    "\n",
    "from getpass import getpass\n",
    "\n",
    "os.environ['OPENAI_API_KEY'] = getpass('Enter your OPENAI_API_KEY key: ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Tw5RzVKl3pUr"
   },
   "source": [
    "**Now you can can try any of the following examples. It is recommended to go through these in sequence, although the order does NOT matter.**\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5rBWNOuXEygx"
   },
   "source": [
    "# Direct interaction with OpenAI LLM\n",
    "\n",
    "In this simple example, we are directly sending a message-sequence to the OpenAI `chatCompletion` API. Note that to have a multi-round converation we have to manually accumulate the dialog.\n",
    "\n",
    "First, import `langroid`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lGY2XyHyD0oJ",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import langroid as lr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = lr.language_models.OpenAIGPT()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bMR6Dani_l9C"
   },
   "source": [
    "We define the LLM model using `OpenAIGPT`; you can optionally pass an `OpenAIGPTConfig` to set the configurations of the OpenAI LLM model.\n",
    "\n",
    "We can also specify the messages that will be sent to instruct the model. `Langroid` supports various roles provided by OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JPul8c4uD1kH"
   },
   "outputs": [],
   "source": [
    "from langroid.language_models import LLMMessage, Role\n",
    "\n",
    "messages = [\n",
    "    LLMMessage(content=\"You are a helpful assistant\",  role=Role.SYSTEM),\n",
    "    LLMMessage(content=\"What is the capital of Ontario?\",  role=Role.USER),\n",
    "]\n",
    "\n",
    "response = llm.chat(messages, max_tokens=200)\n",
    "print(\"LLM response is: \", response.message)\n",
    "\n",
    "# accumulate messages manually\n",
    "\n",
    "messages.append(response.to_LLMMessage())\n",
    "messages.append(LLMMessage(content=\"what about India?\", role=Role.USER))\n",
    "response = llm.chat(messages, max_tokens=200)\n",
    "print(\"LLM response is:\", response.message)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3HzZhfyAfRCQ"
   },
   "source": [
    "The above is a \"raw\" LLM interaction where you have to manage\n",
    "message history. Using an Agent to wrap an LLM, and wrapping an Agent in a Task, we can set up an interactive, multi-round chat much more easily, as we show next."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4TjgYEfK34NZ"
   },
   "source": [
    "## A note on the rest of the examples\n",
    "In the interactive examples below, the conversation loop pauses for human input: in most cases you would hit enter (unless the example requires you to ask a question).\n",
    "The interaction looks much better when run on a terminal,\n",
    "and a notebook is not ideal for these. However we realize a Colab notebook does offer the benefit of having a ready to run environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OaIiDn8zOurc"
   },
   "source": [
    "# Define an agent, set up a task, and run it\n",
    "\n",
    "Say you want to have a multi-round interactive chat with an LLM.\n",
    "\n",
    "`Langroid` simplifies this process. We just need to create a `ChatAgent`, wrap it in a `Task`, and finally run the task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vaXKQWh-SBgH"
   },
   "source": [
    "Note that `Langroid` offers specialized chatting agents such as `DocChatAgent` and `TableChatAgent`, which we will see later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PblTvXC6QZsS"
   },
   "outputs": [],
   "source": [
    "agent = lr.ChatAgent()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LO8O-7vZSoS2"
   },
   "source": [
    "A `ChatAgent` by itself offers 3 standard \"responders\": the LLM, the human User, and the Agent itself (e.g. to handle tool/function-calling by the LLM). To use these responders in an interactive loop, we need to wrap the Agent in a task,\n",
    "and call its `run()` method.\n",
    "\n",
    "As before, a `ChatAgent` can be configured with an optional `ChatAgentConfig` parameter; here, we use the default behavior. This pattern will repeat throughout `Langroid`.\n",
    "\n",
    "A prompt will be displayed after running this task, so you can interact with the `ChatAgent`.\n",
    "\n",
    "Type your questions and the agent will provide the LLM responses. When done, type `q` to exit.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wheZ7NJGO2X4"
   },
   "outputs": [],
   "source": [
    "agent.message_history.clear()\n",
    "task = lr.Task(agent, name=\"Bot\")\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we don't require custom behavior from the `Agent`, this is even simpler: `Task` will use a `ChatAgent` by default."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "B4q8ojyAQ_OV"
   },
   "source": [
    "# Three communicating agents\n",
    "\n",
    "The above example involved a single `ChatAgent`, but in non-trivial applications, we will often find it easier to divide responsibilities among multiple agents, each with different skills and responsibilities.\n",
    "\n",
    "If you attempt to solve these with a single Agent, you would have to keep track of multiple conversation states and loops, and it quickly gets out of hand. Agents offer a way to solve complex tasks in a modular fashion. Moreover, specialized agents can be designed and tested in isolation, and then combined to solve various tasks.\n",
    "\n",
    "`Langroid` streamlines the process of setting up multiple agents and orchestrating their interaction. Here's a toy numerical example (this helps keep token costs low!). Imagine a task where we want to construct a series of numbers using the following rule to transform the current number $n$:\n",
    "- if $n$ is even, the next number is $n/2$\n",
    "- if $n$ is odd, the next number is $3n+1$.\n",
    "\n",
    "We can have 3 agents, each wrapped by a `Task`, which collaborate to produce this sequence.\n",
    "Given the current number $n$,\n",
    "- `repeater_task` simply returns $n$,\n",
    "- `even_task` specializes in handling even numbers, and returns $n/2$ if $n$ is even, else says \"DO-NOT-KNOW\"\n",
    "- `odd_task` specializes in handling odd numbers, and returns $3*n+1$ if $n$ is odd, else says \"DO-NOT-KNOW\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oUbVybSuRIFX"
   },
   "outputs": [],
   "source": [
    "NO_ANSWER = lr.utils.constants.NO_ANSWER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_1aeRre35Ghd"
   },
   "source": [
    "As before, we define chat model that will be used by the agents:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xpqzS-ozUowm"
   },
   "source": [
    "Now, we create the `repeater_task`; note that, as we want to use standard `ChatAgent` behavior, there is no need to configure the `Agent`. The `Task` comprises the following settings:\n",
    "\n",
    "\n",
    "*   **Name**: name of the agent\n",
    "*   **llm_delegate**: whether to delegate control to LLM; conceptually, the \"controlling entity\" is the one \"seeking\" responses to its queries, and has a goal it is aiming to achieve. The \"controlling entity\" is either the LLM or the USER. (Note within a Task there is just one LLM, and all other entities are proxies of the \"User\" entity).\n",
    "*   **single_round**: If true, the task runs until one message by the controller and a subsequent response by the non-controller. If false, runs for the specified number of turns in `run`, or until `done()` is true.\n",
    "* **system_message**: provides instructions to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AI5oFOM2SLgM"
   },
   "outputs": [],
   "source": [
    "repeater_task = lr.Task(\n",
    "    name = \"Repeater\",\n",
    "    system_message=\"\"\"\n",
    "    Your job is to repeat whatever number you receive.\n",
    "    \"\"\",\n",
    "    llm_delegate=True, # LLM takes charge of task\n",
    "    single_round=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zVJ-nbs08Ub5"
   },
   "source": [
    "Now we define our task `even_task`; as before, this task creates its own associated `ChatAgent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pgOFydVqRbPc"
   },
   "outputs": [],
   "source": [
    "even_task = lr.Task(\n",
    "    name = \"EvenHandler\",\n",
    "    system_message=f\"\"\"\n",
    "    You will be given a number.\n",
    "    If it is even, divide by 2 and say the result, nothing else.\n",
    "    If it is odd, say {NO_ANSWER}\n",
    "    \"\"\",\n",
    "    single_round=True,  # task done after 1 step() with valid response\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z-RAYL9S9XTM"
   },
   "source": [
    "Finally, we create the 3rd task `odd_task`; this task again creates an associated `ChatAgent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ukf79CqIScHQ"
   },
   "outputs": [],
   "source": [
    "odd_task = lr.Task(\n",
    "    name = \"OddHandler\",\n",
    "    system_message=f\"\"\"\n",
    "    You will be given a number n.\n",
    "    If it is odd, return (n*3+1), say nothing else.\n",
    "    If it is even, say {NO_ANSWER}\n",
    "    \"\"\",\n",
    "    single_round=True,  # task done after 1 step() with valid response\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PbGYXCpR9uVH"
   },
   "source": [
    "We use `add_sub_task` to orchestrate the collaboration between the agents.  Specifically, `repeater_task` will act as the \"main\", and we add `even_task` and `odd_task` as\n",
    "subtasks. For more details see these [docs](https://langroid.github.io/langroid/quick-start/multi-agent-task-delegation/#task-collaboration-via-sub-tasks).\n",
    "\n",
    "\n",
    "Finally, we kickoff the task with a starting number 3, using `repeater_task.run(\"3\")`.\n",
    "\n",
    "Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0vy77Os8TAas"
   },
   "outputs": [],
   "source": [
    "repeater_task.add_sub_task([even_task, odd_task])\n",
    "repeater_task.set_color_log(enable=False)\n",
    "repeater_task.run(\"3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "t7c3qKvcTwTG"
   },
   "source": [
    "# Simple Tool/Function-calling example\n",
    "\n",
    "Here is a simple numerical example showcasing how `Langroid` supports tools/function-calling. For more details see these [doc pages](https://langroid.github.io/langroid/quick-start/chat-agent-tool/)\n",
    "\n",
    "Say the agent has a secret list of numbers, and we want the LLM to find the smallest number in the list. We want to give the LLM the ability to use a **probe** tool/function which takes a single number `n` as an argument. The tool handler method in the agent returns how many numbers in its list are at most `n`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "htqVMFU2pog6"
   },
   "source": [
    "To use tools/function-calling in `Langroid`, we first **define** the tool as a subclass of `ToolMessage` to specify some details about the tool (e.g., name and parameters) and when it can be used/triggered:\n",
    "* **request**: is the name of the tool/function, as well as the name of the Agent method that \"handles\" the tool.\n",
    "* **purpose**: general description to give hints to LLM when this tool can be used\n",
    "* **number**: is a function-argument for the `probe` tool and its type is `int`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3eIidMmZUgir"
   },
   "outputs": [],
   "source": [
    "class ProbeTool(lr.agent.ToolMessage):\n",
    "  request: str = \"probe\"\n",
    "  purpose: str = \"\"\"\n",
    "        To find how many numbers in my list are less than or equal to\n",
    "        the <number> you specify.\n",
    "        \"\"\" # note  <number> corresponds to the name of the tool's argument/parameter\n",
    "  number: int"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vBAw8kxkqa42"
   },
   "source": [
    "Next, we create an agent `SpyGameAgent`, with a special method `probe` to handle the `probe` tool/function.\n",
    "Notice the argument of the `probe` method is an instance of the class `ProbeTool` that we created in the previous step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "einGIzjQUyX7"
   },
   "outputs": [],
   "source": [
    "class SpyGameAgent(lr.ChatAgent):\n",
    "  def __init__(self, config: lr.ChatAgentConfig = lr.ChatAgentConfig()):\n",
    "    super().__init__(config)\n",
    "    self.numbers = [3, 4, 8, 11, 15, 25, 40, 80, 90] # agent's secret list\n",
    "\n",
    "  def probe(self, msg: ProbeTool) -> str:\n",
    "    # return how many numbers in self.numbers are less or equal to msg.number\n",
    "    return str(len([n for n in self.numbers if n <= msg.number]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pG2MJJN4yCHb"
   },
   "source": [
    "Finally, we instantiate the `SpyGameAgent` as an object `spy_game_agent`, and \"associate\" the `probe` tool with this agent, using the `enable_message` method of the `ChatAgent`.  We then wrap the `spy_game_agent` in a `Task` object, with instructions (`system_message`) on what it should aim for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OM6Lk3uWVOCB"
   },
   "outputs": [],
   "source": [
    "spy_game_agent = SpyGameAgent()\n",
    "\n",
    "spy_game_agent.enable_message(ProbeTool)\n",
    "\n",
    "task = lr.Task(\n",
    "        spy_game_agent,\n",
    "        name=\"Spy\",\n",
    "        system_message=\"\"\"\n",
    "            I have a list of numbers between 1 and 20.\n",
    "            Your job is to find the smallest of them.\n",
    "            To help with this, you can give me a number and I will\n",
    "            tell you how many of my numbers are equal or less than your number.\n",
    "            Once you have found the smallest number,\n",
    "            you can say DONE and report your answer.\n",
    "        \"\"\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aQ93n0kA_rM0"
   },
   "source": [
    "Now run the task.\n",
    "\n",
    "Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xAVoOcgsNWOv"
   },
   "outputs": [],
   "source": [
    "spy_game_agent.message_history.clear()\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qmeh3zJTZeL1"
   },
   "source": [
    "# Chat with documents (file paths, URLs, etc)\n",
    "\n",
    "In the previous examples, the Agents did not use any external documents. In this example, we we set up an Agent that supports \"chatting\" with documents. Specifically, we use the `DocChatAgent` class to ask questions about a set of URLs.\n",
    "The `DocChatAgent` first ingests the contents of the websites specified by the URLs by chunking, embedding and indexing them into a vector database (`qdrant` by default). We then wrap the agent in a task and run it interactively.\n",
    "The user can ask questions and the LLM of the agent returns answers using Retrieval Augment Generation, with Evidence Citation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langroid.agent.special import DocChatAgent, DocChatAgentConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w1ZcFRJu5K5D"
   },
   "source": [
    "Now we define the configuration of the `DocChatAgent`. The configurations include the path to access the documents, chat model settings, and vector-DB settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tVEHeg9jZpl7"
   },
   "outputs": [],
   "source": [
    "config = DocChatAgentConfig(\n",
    "    doc_paths = [\n",
    "        \"https://en.wikipedia.org/wiki/Language_model\",\n",
    "        \"https://en.wikipedia.org/wiki/N-gram_language_model\",\n",
    "    ],\n",
    "    vecdb=lr.vector_store.QdrantDBConfig(\n",
    "        collection_name=\"docqa-chat-multi-extract\",\n",
    "        storage_path=\".qdrant/test2/\", # CHANGE THIS PATH IF YOU GET AN ERROR WHEN RE-RUNNING THE CELL\n",
    "    ),\n",
    ")\n",
    "\n",
    "agent = DocChatAgent(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UZyTx3vQN_I4"
   },
   "source": [
    "As before, we wrap the agent in a task, and run it.\n",
    "\n",
    "Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OMFDp8WxaAI3"
   },
   "outputs": [],
   "source": [
    "agent.message_history.clear()\n",
    "task = lr.Task(agent)\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kZqX1J6qWhHk"
   },
   "source": [
    "# Tool/Function-calling to extract structured information from text\n",
    "\n",
    "Let's combine multi-agent interaction, Retrieval-Augmented Generation, and tools/function-calling, for a more realistic example. Suppose you want an agent to extract the key terms of a lease, from a lease document, as a nested JSON structure.\n",
    "This can be accomplished by instructing the LLM to use a specific tool.\n",
    "\n",
    "To simplify the solution, we separate the skills/responsibilities into two different Agents:\n",
    "- `LeaseExtractorAgent` has no access to the lease, and is responsible for gathering the key terms into a specific structured form\n",
    "- `DocChatAgent` has access to the lease and answers specific questions it receives from the `LeaseExtractorAgent`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from pydantic import BaseModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d-CkI7wk5A8l"
   },
   "source": [
    "Next, we define the desired structure of the lease information via Pydantic models. The desired format is a nested JSON structure, which maps to a nested class structure:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0p8iveEcX_1E"
   },
   "outputs": [],
   "source": [
    "class LeasePeriod(BaseModel):\n",
    "    start_date: str\n",
    "    end_date: str\n",
    "\n",
    "class LeaseFinancials(BaseModel):\n",
    "    monthly_rent: str\n",
    "    deposit: str\n",
    "\n",
    "class Lease(BaseModel):\n",
    "    \"\"\"\n",
    "    Various lease terms.\n",
    "    Nested fields to make this more interesting/realistic\n",
    "    \"\"\"\n",
    "\n",
    "    period: LeasePeriod\n",
    "    financials: LeaseFinancials\n",
    "    address: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WcUQylLu5HFh"
   },
   "source": [
    "We then define the `LeaseMessage` tool as a subclass of Langroid's `ToolMessage`. The `LeaseMessage` class has a\n",
    "required argument `terms` of type `Lease`. The `classmethod` named `examples` is used to generate $k$-shot examples for the LLM when instructing it to extract information in the desired structured form (see a later cell below).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XFVCpL8jW7C7"
   },
   "outputs": [],
   "source": [
    "class LeaseMessage(lr.agent.ToolMessage):\n",
    "    request: str = \"lease_info\" # maps to method of LeaseExtractorAgent\n",
    "    purpose: str = \"\"\"\n",
    "        Collect information about a Commercial Lease.\n",
    "        \"\"\"\n",
    "    terms: Lease\n",
    "\n",
    "    @classmethod\n",
    "    def examples(cls) -> List[\"LeaseMessage\"]:\n",
    "        return [\n",
    "            cls(\n",
    "                terms=Lease(\n",
    "                    period=LeasePeriod(start_date=\"2021-01-01\", end_date=\"2021-12-31\"),\n",
    "                    financials=LeaseFinancials(monthly_rent=\"$1000\", deposit=\"$1000\"),\n",
    "                    address=\"123 Main St, San Francisco, CA 94105\",\n",
    "                ),\n",
    "                result=\"\",\n",
    "            ),\n",
    "        ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lu03iGEaW0Ur"
   },
   "source": [
    "Next we define the `LeaseExtractorAgent` and add a method `least_info` to handle the tool/function-call `lease_info` defined in the tool `LeaseMessage`. In this case the handling is trivial: if the method receives a valid object of class `LeaseMessage`, it declares \"success\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZlZ0UtqEXGdz"
   },
   "outputs": [],
   "source": [
    "class LeaseExtractorAgent(lr.ChatAgent):\n",
    "    def __init__(self, config: lr.ChatAgentConfig = lr.ChatAgentConfig()):\n",
    "        super().__init__(config)\n",
    "\n",
    "    def lease_info(self, message: LeaseMessage) -> str:\n",
    "        print(\n",
    "            f\"\"\"\n",
    "        DONE! Successfully extracted Lease Info:\n",
    "        {message.terms}\n",
    "        \"\"\"\n",
    "        )\n",
    "        return json.dumps(message.terms.dict())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HtHXm0-UuBRU"
   },
   "outputs": [],
   "source": [
    "# Obtain the lease.txt document that we want to parsed\n",
    "!wget https://github.com/langroid/langroid-examples/blob/main/examples/docqa/lease.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mdkxu37f7JuK"
   },
   "source": [
    "Next, set up an instance of `DocChatAgent`, point it to the lease document, equip it with a vector database, and instructions on how to answer questions based on extracts retrieved from the vector-store.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d8ynRPp6XUvV"
   },
   "outputs": [],
   "source": [
    "doc_agent = DocChatAgent(\n",
    "        DocChatAgentConfig(\n",
    "            doc_paths = [\"lease.txt\"],\n",
    "            vecdb=lr.vector_store.QdrantDBConfig(\n",
    "                collection_name=\"docqa-chat-multi-extract\",\n",
    "                storage_path=\".data1/data1/\", # CHANGE PATH IF ERROR\n",
    "              ),\n",
    "            summarize_prompt= f\"\"\"\n",
    "                Use the provided extracts to answer the question.\n",
    "                If there's not enough information, respond with {NO_ANSWER}. Use only the\n",
    "                information in these extracts, even if your answer is factually incorrect,\n",
    "                and even if the answer contradicts other parts of the document. The only\n",
    "                important thing is that your answer is consistent with and supported by the\n",
    "                extracts. Compose your complete answer and cite all supporting sources on a\n",
    "                separate separate line as \"EXTRACTS:\".\n",
    "                Show each EXTRACT very COMPACTLY, i.e. only show a few words from\n",
    "                the start and end of the extract, for example:\n",
    "                EXTRACT: \"The world war started in ... Germany Surrendered\"\n",
    "                {{extracts}}\n",
    "                {{question}}\n",
    "                Answer:\n",
    "            \"\"\"\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DISYqYVr8BQx"
   },
   "source": [
    "Next we wrap the `doc_agent` into a Task, with instructions on its role.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NBLp0AEy74N-"
   },
   "outputs": [],
   "source": [
    "doc_task = lr.Task(\n",
    "    doc_agent,\n",
    "    name=\"DocAgent\",\n",
    "    llm_delegate=False,\n",
    "    single_round=True,\n",
    "    system_message=\"\"\"You are an expert on Commercial Leases.\n",
    "    You will receive various questions about a Commercial\n",
    "    Lease contract, and your job is to answer them concisely in at most 2 sentences.\n",
    "    Please SUPPORT your answer with an actual EXTRACT from the lease,\n",
    "    showing only a few words from the  START and END of the extract.\n",
    "    \"\"\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZhRfOPiE9W-n"
   },
   "source": [
    "Finally, we instantiate the `lease_extractor_agent`, enable it to use and handle the `LeaseMessage` tool. Then we wrap the `lease_extractor_agent` into a Task, instructing it to gather information in the desired format, by asking questions one at a time. Note how the instruction contains `LeaseMessage.usage_example()`: this example is constructed from the `examples` classmethod above when the `LeaseMessage` was defined.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VL2RYZUX7393"
   },
   "outputs": [],
   "source": [
    "lease_extractor_agent = LeaseExtractorAgent()\n",
    "\n",
    "lease_extractor_agent.enable_message(\n",
    "    LeaseMessage,\n",
    "    use=True,\n",
    "    handle=True,\n",
    "    force=False,\n",
    ")\n",
    "\n",
    "lease_task = lr.Task(\n",
    "    lease_extractor_agent,\n",
    "    name=\"LeaseExtractorAgent\",\n",
    "    llm_delegate=True,\n",
    "    single_round=False,\n",
    "    system_message=f\"\"\"\n",
    "    You have to collect some information about a Commercial Lease, but you do not\n",
    "    have access to the lease itself.\n",
    "    You can ask me questions about the lease, ONE AT A TIME, I will answer each\n",
    "    question. You only need to collect info corresponding to the fields in this\n",
    "    example:\n",
    "    {LeaseMessage.usage_example()}\n",
    "    If some info cannot be found, fill in {NO_ANSWER}.\n",
    "    When you have collected this info, present it to me using the\n",
    "    'lease_info' function/tool.\n",
    "    \"\"\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XBdAcarpwnYU"
   },
   "source": [
    "Finally, we set up the `doc_task` as a subtask of the `lease_task` so that the `doc_agent` can respond to questions from the `lease_extractor_agent`.\n",
    " Now, the `lease_extractor_agent` will be asking questions about the lease and `doc_task` will provide the answers, citing evidence extracted from the lease. Once `lease_extractor_agent` collects all the terms of the lease as instructed, it will use the tool `LeaseMessage` to return this information.\n",
    "\n",
    " The next cell runs the `lease_task`. Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1Lt9FhjTFlh2"
   },
   "outputs": [],
   "source": [
    "lease_extractor_agent.message_history.clear()\n",
    "lease_task.add_sub_task(doc_task)\n",
    "lease_task.set_color_log(enable=False)\n",
    "lease_task.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uWcH7HoFc-am"
   },
   "source": [
    "# Chat with tabular data (file paths, URLs, dataframes)\n",
    "\n",
    "Here is how `Langroid's` `TableChatAgent` can be used to chat with tabular data, which can be specified as a URL, file path or Pandas dataframe.\n",
    "\n",
    "The Agent's LLM generates Pandas code to answer the query, via function-calling (or tool/plugin), and the Agent's function-handling method executes the code and returns the answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jS06lshgdBv7"
   },
   "outputs": [],
   "source": [
    "from langroid.agent.special.table_chat_agent import TableChatAgent, TableChatAgentConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tS8pDD2MdgPi"
   },
   "source": [
    "Set up a `TableChatAgent` for a data file, URL or dataframe (Ensure the data table has a header row; the delimiter/separator is auto-detected):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qDbA4VGmda10"
   },
   "outputs": [],
   "source": [
    "dataset =  \"https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv\"\n",
    "# or dataset = \"/path/to/my/data.csv\"\n",
    "# or dataset = pd.read_csv(\"/path/to/my/data.csv\")\n",
    "\n",
    "agent = TableChatAgent(\n",
    "    config=TableChatAgentConfig(\n",
    "        data=dataset,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "em2hcy2Qd67T"
   },
   "source": [
    "Now, let's set up a task and run it in an interactive loop with the user:\n",
    "Based on `dataset`, you can ask the following question in the prompt:\n",
    "\n",
    "```\n",
    "What is the average alcohol content of wines with a quality rating above 7?\n",
    "```\n",
    "\n",
    "Remember to keep hitting enter when it's the human's turn, and hit \"q\" to end the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nhDMm5ndd93W"
   },
   "outputs": [],
   "source": [
    "agent.message_history.clear()\n",
    "task = lr.Task(agent, name=\"DataAssistant\")\n",
    "task.set_color_log(enable=False)\n",
    "task.run()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "include_colab_link": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
